import unittest

import pandas as pd
from matplotlib import pyplot as plt
from sklearn.linear_model import LinearRegression
import numpy as np
from sklearn.metrics import mean_squared_error,r2_score

"""
    线性回归模型预测
    公众号：干货食堂
    Date : 2023-11-19 21:08:12
"""
class MyTestCase(unittest.TestCase):

    def test(self):
        x = np.arange(0, 20)
        # 生成随机数
        y = np.random.randint(0, 20, size=20)
        # 设置图片大小
        plt.figure(figsize=(5, 5))
        plt.rcParams['font.family'] = 'Microsoft YaHei'
        plt.title('散点图')
        plt.xlabel('x 轴')
        plt.ylabel('y 轴')
        # plt.plot(x, y, 'ob')
        plt.scatter(x, y)
        plt.show()

    def test2(self):
        pass

    def test_something(self):
        data = pd.read_csv('usa_housing_price.csv')

        fig = plt.figure(figsize=(10, 10))
        fig1 = plt.subplot(231)
        plt.scatter(data.loc[:, 'Avg. Area Income'], data.loc[:, 'Price'])
        plt.title('Price VS Income')

        fig2 = plt.subplot(232)
        plt.scatter(data.loc[:, 'Avg. Area House Age'], data.loc[:, 'Price'])
        plt.title('Price VS House Age')

        fig3 = plt.subplot(233)
        plt.scatter(data.loc[:, 'Avg. Area Number of Rooms'], data.loc[:, 'Price'])
        plt.title('Price VS Number of Rooms')

        fig4 = plt.subplot(234)
        plt.scatter(data.loc[:, 'Area Population'], data.loc[:, 'Price'])
        plt.title('Price VS Area Population')

        fig5 = plt.subplot(235)
        plt.scatter(data.loc[:, 'size'], data.loc[:, 'Price'])
        plt.title('Price VS size')
        # plt.show()


        X = data.loc[:, 'size']
        y = data.loc[:, 'Price']
        X = np.array(X).reshape(-1, 1)


        LR1 = LinearRegression()
        LR1.fit(X, y)
        y_predict_1 = LR1.predict(X)

        fig6 = plt.figure(figsize=(8, 5))
        plt.scatter(X, y)
        plt.plot(X, y_predict_1, 'r')
        # plt.show()
        # 检验预测正确性
        mean_squared_error_1 = mean_squared_error(y, y_predict_1)
        r2_score_1 = r2_score(y, y_predict_1)
        print(mean_squared_error_1, r2_score_1)

        # 多维度预测
        X_multi = data.drop(['Price'], axis=1)
        LR_multi = LinearRegression()
        # train the model
        LR_multi.fit(X_multi, y)
        y_predict_multi = LR_multi.predict(X_multi)
        print(y_predict_multi)

        mean_squared_error_multi = mean_squared_error(y, y_predict_multi)
        r2_score_multi = r2_score(y, y_predict_multi)
        print(mean_squared_error_multi, r2_score_multi)

        fig7 = plt.figure(figsize=(8, 5))
        plt.scatter(y, y_predict_multi)
        # plt.show()

        X_test = [65000, 5, 5, 30000, 200]
        X_test = np.array(X_test).reshape(1, -1)
        y_test_predict = LR_multi.predict(X_test)
        print(y_test_predict)